//% color="#6495ED" iconWidth=50 iconHeight=40
namespace MMEdu{
    //% block="训练模型" blockType="tag"
    export function tagtest() {}
    //% block="初始化MMEdu模块,选择[MMDEU]模块,设置[MODEL]模型" blockType="command" 
    //% MMDEU.shadow="dropdown"   MMDEU.options="MMDEU"
    //% MODEL.shadow="dropdown"   MODEL.options="MODEL"
    export function MMEdu_Init(parameter: any, block: any){
        let MMDEU = parameter.MMDEU.code
        let model = parameter.MODEL.code
      
        Generator.addImport(`from MMEdu import ${MMDEU} as mmdeu_cls_det`)
        Generator.addImport("import os")
        Generator.addImport("import time")
        Generator.addCode(`mmdeu_model = mmdeu_cls_det("${model}")`)

}
    //% block="指定数据集中的类别数量[NUMBER]" blockType="command"
    //% NUMBER.shadow="normal"   NUMBER.defl="3"
    export function MMEdu_Init_ncl(parameter: any, block: any){
        let number = parameter.NUMBER.code
        Generator.addCode(`mmdeu_model.num_classes = ${number}`)
    }


    //% block="从指定数据集路径中加载数据[PATH]" blockType="command"
    //% PATH.shadow="string"   PATH.defl="/path/to/dataset"
    export function MMEdu_Init_loaddata(parameter: any, block: any){
        let path = parameter.PATH.code

        Generator.addCode(`mmdeu_model.load_dataset(path=${path})`)
    }

    //% block="设置模型的保存路径[PATH]" blockType="command"
    //% PATH.shadow="string"   PATH.defl="/path/to/cls"
    export function MMEdu_Init_sf(parameter: any, block: any){
        let path = parameter.PATH.code
        Generator.addCode(`
if not os.path.exists(${path}):
    print("The folder does not exist,created automatically")
    os.makedirs( ${path})
 
mmdeu_model.save_fold = ${path}`)
    }

    //% block="训练模型，设置参数训练[NUM]次, 选择[OPT]优化器,学习率[LR],预训练权重文件加载路径[CP]" blockType="command"
    //% NUM.shadow="normal" NUM.defl="10"
    //% OPT.shadow="dropdown" OPT.options="OPT"
    //% LR.shadow="normal" LR.defl="0.01"
    //% CP.shadow="normal" CP.defl=None
    export function MMEdu_Init_train(parameter: any, block: any) {
        let NUM = parameter.NUM.code
        let OPT = parameter.OPT.code
        let LR = parameter.LR.code
        let CP = parameter.CP.code
        if (`${CP}` === "None" ) {
            CP = "None"
        }else { 
            CP = `r"${CP}"`}
        Generator.addCode(`
try:
    mmdeu_model.train(epochs=${NUM}, validate=True, optimizer="${OPT}", lr=${LR}, checkpoint=${CP})
except:
    pass
`) 
}

    //% block="模型转换" blockType="tag"
    export function tagtest1() {}
    //% block="设置模型转换参数[CONFIG]为[VAL]" blockType="command"
    //% CONFIG.shadow="dropdown"   CONFIG.options="CONFIG"
    //% VAL.shadow="normal"   VAL.defl=""
    export function MMEdu_Init_mc(parameter: any, block: any){
        
        let CONFIG = parameter.CONFIG.code
        let VAL = parameter.VAL.code
        Generator.addImport("from MMEdu import MMClassification as cls")
        Generator.addImport("import os\nimport glob\npara_model_convert = {}")
    
 
        if (`${CONFIG}` === "checkpoint") {
            Generator.addCode(`
CONFIG = "${CONFIG}"  
VAL = r"${VAL}" 
if CONFIG == "checkpoint":
    pattern = os.path.join(VAL, 'best_accuracy*.pth')
    matching_files = glob.glob(pattern)
    if matching_files:
        para_model_convert[CONFIG] = matching_files[0]  
    else:
        print("No matching .pth file found in the specified directory.")
`)
        }
        else {

            Generator.addCode(`
output_directory = os.path.dirname("${VAL}")
if not os.path.exists(output_directory):
    print("The folder does not exist,created automatically")
    os.makedirs(output_directory)
 
para_model_convert["${CONFIG}"] = "${VAL}"`)
        }

    }
    //% block="完成模型转换参数设置，进行一次模型转换" blockType="command"
    export function MMEdu_Init_(parameter: any, block: any){
        Generator.addCode(`mmdeu_model.convert(**para_model_convert)`)
    }
}